#!/usr/bin/env python3
'''
已验证，验证不通过
1. 无法完成一局游戏
2. play的过程中容易停止不动好几秒，然后再继续
3. 已经涂过颜色的墙壁依旧会重走

训练记录：
1029：训练分数最高达到24分（高低分震荡），测试分数最高达到97.8分（接近今天训练结束时产生），总体应该是上涨的，明天继续训练看看
1030：训练分数达到62分左右（震荡幅度不大），测试分数达到112.4，继续训练看看
1031:训练分数达到108（震荡幅度不大），测试分数达到170.4，继续训练看看
1101：训练分数达到161（震荡幅度不大），测试分数达到263.4，继续训练看看
1102:训练分数达到207分（振幅貌似有点大），测试分数达到299.4，继续训练看看
1103：训练分数196（对比昨天分数貌似降低），测试分数317，继续训练
以上都没加入超参数调整，后续无提升可先不加入，让图形追上再说
1104: 训练分数209分，测试分数达到317.6分，暂时停止训练，训练卷积达到该指标
20250219: 完成代码调整，重新进行挑战,训练分数53.870，测试分数63.4，继续训练
20250220: 继续开始训练，训练分数59.5，测试分数71.6，继续训练
20250221: 继续开始训练，训练分数66.740，测试分数87.2，继续训练
20250222: 继续开始训练，训练分数78.610，测试分数88分，继续训练
20250224: 继续开始训练，训练分数72.4，测试分数90.8，继续训练
20250225: 继续开始训练，训练分数77.120，测试分数92.8，继续训练
20250226: 继续开始训练，训练分数84.530，测试分数94分，继续训练
20250227: 继续开始训练，开放学习率调度器，训练分数79.740，测试分数94.8，继续训练，查看调度器和学习率
20250228：继续开始训练，学习率1e-4，训练分数83.400，测试分数95.6，继续训练
20250301:继续开始训练，当前学习率： 5.904900000000002e-05，训练分数87.310，测试分数97.6，继续训练，学习率配置失败，需要检查
20250302:o继续开始训练，学习率： 5.904900000000002e-05，训练分数82.980，测试分数99分，继续训练，对比恢复原来的代码，为什么训练的进度这么慢
20250303:重新恢复学习率当前学习率： 0.0001，测试分数99分，训练分数78.690，继续训练一天并恢复对比原来的模型
todo 对比恢复原来的代码，为什么训练的进度这么慢
20250304：当前学习率： 0.0001，测试分数99分，训练分数未知，继续训练
20250305:当前学习率： 9e-05，测试分数130.8分，训练分数105.560，继续训练，关闭学习率调度器
发现了大量的缩放过大奖励的日志输出，说明分数相差太大，就是因为奖励的缩放过大，训练代码整体没有问题
20250306:当前学习率： 6.561000000000002e-05，测试分数135.4分，训练分数未知，继续训练一天
20250307:当前学习率： 6.561000000000002e-05，测试分数137.6，训练分数116.970，继续训练一天
20250308:当前学习率： 6.561000000000002e-05，测试分数140.4，训练分数120.970，继续训练一天，感觉训练进度有点慢，是不是学习率太低了
20250309:当前学习率： 6.561000000000002e-05，测试分数144分，训练分数124.10，继续训练
200250310：当前学习率： 6.561000000000002e-05，测试分数144分，训练分数113.940，停止训练，对比原先的代码，重新训练，总结不同点
'''
import gymnasium as gym
import ptan
import numpy as np
import argparse
from tensorboardX import SummaryWriter
import os

import torch
import torch.nn as nn
import torch.nn.utils as nn_utils
import torch.nn.functional as F
import torch.optim as optim
from collections import deque
import cv2
from lib import common
import ale_py


from lib import common

gym.register_envs(ale_py)
GAMMA = 0.99
LEARNING_RATE = 1e-4
ENTROPY_BETA = 0.01
BATCH_SIZE = 128
NUM_ENVS = 50
SAVE_ITERS = 100

REWARD_STEPS = 4
CLIP_GRAD = 0.5

class AtariA2C(nn.Module):
    def __init__(self, obs_size, n_actions):
        super(AtariA2C, self).__init__()

        self.conv = nn.Sequential(
            nn.Linear(obs_size, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 768),
            nn.ReLU(),
            nn.Linear(768, 1024),  # New layer
            nn.ReLU(),
        )

        self.policy = nn.Sequential(
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, n_actions)
        )

        self.value = nn.Sequential(
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

    def forward(self, x):
        linear_out = self.conv(x.float())
        return self.policy(linear_out), self.value(linear_out)
    

class RewardPenaltyWrapper(gym.Wrapper):
    def __init__(self, env, frame_penalty=-0.1, life_loss_penalty=-10):
        super(RewardPenaltyWrapper, self).__init__(env)
        self.frame_penalty = frame_penalty
        self.life_loss_penalty = life_loss_penalty
        self.previous_lives = 0

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        self.previous_lives = info.get('lives', 0)  # 初始生命值
        return obs, info

    def step(self, action):
        obs, reward, done, truncated, info = self.env.step(action)
        
        # 每帧惩罚
        # reward += self.frame_penalty
        
        # 处理生命减少时的惩罚
        current_lives = info.get('lives', self.previous_lives)
        if current_lives < self.previous_lives:
            reward += self.life_loss_penalty
            self.previous_lives = current_lives
        
        return obs, reward, done, truncated, info
    

def wrap_dqn(env, stack_frames=4, episodic_life=True, reward_clipping=True):
    if episodic_life:
        # 将多条生命的游戏模拟成单条生命ActorCriticAgent
        env = ptan.common.wrappers.EpisodicLifeEnv(env)
    # 增强初始化
    env = ptan.common.wrappers.NoopResetEnv(env, noop_max=30)

    if 'FIRE' in env.unwrapped.get_action_meanings():
        env = ptan.common.wrappers.FireResetEnv(env)
    env = common.BreakoutPenaltyWrapper(env)
    return env


def unpack_batch(batch, net, device='cpu'):
    """
    Convert batch into training tensors
    :param batch:
    :param net:
    :return: states variable, actions tensor, reference values variable
    """
    states = []
    actions = []
    rewards = []
    not_done_idx = [] # 非结束的游戏数据索引，该索引记录对应batch，states，actions，rewards
    last_states = [] # 记录采样中的执行动作后的状态，仅记录游戏非结束状态下的索引
    for idx, exp in enumerate(batch):
        states.append(np.array(exp.state, copy=False))
        actions.append(int(exp.action))
        rewards.append(exp.reward)
        if exp.last_state is not None:
            not_done_idx.append(idx)
            last_states.append(np.array(exp.last_state, copy=False))
    states_v = torch.FloatTensor(np.array(states, copy=False)).to(device)
    actions_t = torch.LongTensor(actions).to(device)
    rewards_np = np.array(rewards, dtype=np.float32)
    if not_done_idx:
        last_states_v = torch.FloatTensor(np.array(last_states, copy=False)).to(device)
        last_vals_v = net(last_states_v)[1]
        last_vals_np = last_vals_v.data.cpu().numpy()[:, 0]
        rewards_np[not_done_idx] += GAMMA ** REWARD_STEPS * last_vals_np

    ref_vals_v = torch.FloatTensor(rewards_np).to(device)
    return states_v, actions_t, ref_vals_v

def test_model(env, net, device, episodes=5):
    total_reward = 0.0
    for _ in range(episodes):
        obs, _ = env.reset()
        while True:
            obs_v = ptan.agent.default_states_preprocessor([obs]).to(device)
            logits_v, _ = net(obs_v)
            probs_v = F.softmax(logits_v, dim=1)
            probs = probs_v.data.cpu().numpy()
            action = np.argmax(probs)
            obs, reward, done, trunc, _ = env.step(action)
            total_reward += reward
            if done or trunc:
                break
    return total_reward / episodes


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=True, action="store_true", help="Enable cuda")
    parser.add_argument("-n", "--name", default="ppo_carracing_linear", required=False, help="Name of the run")
    args = parser.parse_args()
    device = torch.device("cuda" if args.cuda else "cpu")

    save_path = os.path.join("saves", "a2c-linear-" + args.name)
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    envs = [wrap_dqn(gym.make("ALE/Amidar-v5", frameskip=4, repeat_action_probability=0.0, obs_type="ram")) for _ in range(NUM_ENVS)]
    test_env = wrap_dqn(gym.make("ALE/Amidar-v5", frameskip=4, repeat_action_probability=0.0, obs_type="ram"))
    writer = SummaryWriter(comment="-alien-a2c-linear-" + args.name)

    net = AtariA2C(envs[0].observation_space.shape[0], envs[0].action_space.n).to(device)
    print(net)

    agent = ptan.agent.PolicyAgent(lambda x: net(x)[0], apply_softmax=True, device=device)
    exp_source = ptan.experience.ExperienceSourceFirstLast(envs, agent, gamma=GAMMA, steps_count=REWARD_STEPS)
    optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE, eps=1e-3)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50000, gamma=0.9)

    start_idx = 0
     # 增加加载模型的代码
    if os.path.exists(save_path) and len(os.listdir(save_path)) > 0:
        # 增加加载模型的代码
        checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(save_path)),
                             key=lambda x: int(x.split('_')[2].split('.')[0]))
        checkpoint = torch.load(os.path.join(save_path, checkpoints[-1]), map_location=device, weights_only=False)
        net.load_state_dict(checkpoint['net'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_idx = checkpoint['start_idx']
        scheduler.load_state_dict(checkpoint['scheduler'])
        print("加载模型成功")
        # 打印学习率
        print(f"optimizer学习率：", optimizer.param_groups[0]['lr'])
        # 手动设置学习率1e-4
        # for param_group in optimizer.param_groups:
        #     param_group['lr'] = 1e-4
        # scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=100000, gamma=0.9)
        print(f"当前学习率：", scheduler.get_last_lr()[0])

    batch = []
    best_reward = 0

    with common.RewardTracker(writer, stop_reward=10000) as tracker:
        with ptan.common.utils.TBMeanTracker(writer, batch_size=10) as tb_tracker:
            for step_idx, exp in enumerate(exp_source):
                batch.append(exp)

                new_rewards = exp_source.pop_total_rewards()
                if new_rewards:
                    if tracker.reward(new_rewards[0], step_idx + start_idx):
                        break

                if len(batch) < BATCH_SIZE:
                    continue

                states_v, actions_t, vals_ref_v = unpack_batch(batch, net, device=device)
                batch.clear()

                optimizer.zero_grad()
                logits_v, value_v = net(states_v)
                loss_value_v = F.mse_loss(value_v.squeeze(-1), vals_ref_v)

                log_prob_v = F.log_softmax(logits_v, dim=1)
                adv_v = vals_ref_v - value_v.squeeze(-1).detach()
                log_prob_actions_v = adv_v * log_prob_v[range(BATCH_SIZE), actions_t]
                loss_policy_v = -log_prob_actions_v.mean()

                prob_v = F.softmax(logits_v, dim=1)
                entropy_loss_v = ENTROPY_BETA * (prob_v * log_prob_v).sum(dim=1).mean()

                loss_policy_v.backward(retain_graph=True)
                grads = np.concatenate([p.grad.data.cpu().numpy().flatten()
                                        for p in net.parameters()
                                        if p.grad is not None])

                loss_v = entropy_loss_v + loss_value_v
                loss_v.backward()
                nn_utils.clip_grad_norm_(net.parameters(), CLIP_GRAD)
                optimizer.step()
                # scheduler.step()
                start_idx += 1
                loss_v += loss_policy_v

                if start_idx % 200 == 0:
                    test_reward = test_model(test_env, net, device=device, episodes=5)
                    print(f"Test reward: {test_reward:.2f}")
                    common.save_best_model(test_reward, net.state_dict(), save_path, "a2c-gray-best", keep_best=10)

                if start_idx % SAVE_ITERS == 0:           
                    checkpoint = {
                        "net": net.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "start_idx": start_idx + step_idx,
                        "scheduler": scheduler.state_dict()
                    }
                    common.save_checkpoints(start_idx + step_idx, checkpoint, save_path, "acktr", keep_last=5)

                tb_tracker.track("advantage",       adv_v, step_idx + start_idx)
                tb_tracker.track("values",          value_v, step_idx + start_idx)
                tb_tracker.track("batch_rewards",   vals_ref_v, step_idx + start_idx)
                tb_tracker.track("loss_entropy",    entropy_loss_v, step_idx + start_idx)
                tb_tracker.track("loss_policy",     loss_policy_v, step_idx + start_idx)
                tb_tracker.track("loss_value",      loss_value_v, step_idx + start_idx)
                tb_tracker.track("loss_total",      loss_v, step_idx + start_idx)
                tb_tracker.track("grad_l2",         np.sqrt(np.mean(np.square(grads))), step_idx + start_idx)
                tb_tracker.track("grad_max",        np.max(np.abs(grads)), step_idx + start_idx)
                tb_tracker.track("grad_var",        np.var(grads), step_idx + start_idx)